package com.xyzboom.compiler.compiler.lr

import com.xyzboom.compiler.compiler.grammar.*
import com.xyzboom.compiler.compiler.grammar.Terminator.Companion.EOF
import com.xyzboom.compiler.compiler.semantic.SemGrammar
import com.xyzboom.compiler.fa.DFA
import com.xyzboom.compiler.fa.EdgeType.*
import com.xyzboom.compiler.fa.FAEdge
import org.jetbrains.annotations.TestOnly

const val PRODUCTION_NAME = "Prod"
const val PRODUCTION_RIGHT_NAME = "Prod_right"
const val GRAMMAR_NAME = "Grammar"
const val SEM_GRAMMAR_NAME = "SemGrammar"
const val SEM_GLOBAL_DEF = "GlobalDef"
const val SEM_EXPRS = "SemExprs"
const val IDS_NAME = "Ids"
const val ID_NAME = "id"
val EX_START_NON_TERMINATOR = NonTerminator("S'")

internal class ProdGroup(
    //分别存储产生式的索引，规约标记的索引
    val reduceMap: HashMap<Pair<Int, Int>, HashSet<Terminator>>,
    var state: Int,
) {
    val prods get() = reduceMap.keys
    override fun equals(other: Any?): Boolean {
        if (this === other) return true
        if (other !is ProdGroup) return false

        if (reduceMap != other.reduceMap) return false

        return true
    }

    override fun hashCode(): Int {
        return reduceMap.hashCode()
    }
}

private operator fun ArrayList<Production>.get(index: Pair<Int, Int>): ISymbol =
    get(index.first)[index.second]

private fun first(
    symbol: NonTerminator, prods: ArrayList<Production>,
    firstMap: HashMap<NonTerminator, HashSet<Terminator>>,
): HashSet<Terminator> =
    firstMap[symbol] ?: run {
        val result = HashSet<Terminator>().apply {
            for (prod in prods) {
                if (prod.left == symbol) {
                    if (prod.right.size == 0) {//右部为空，加入空串
                        add(EOF)
                    } else {
                        var hasEOF = true
                        for (r in prod.right) {//遍历右部，直到找到一非空的first集
                            if (r.isTerminator) {
                                add(r as Terminator)
                                break
                            } else {
                                val firstR = first(r as NonTerminator, prods, firstMap)
                                if (firstR.contains(EOF)) {
                                    addAll(firstR.apply { remove(EOF) })
                                } else {
                                    addAll(firstR)
                                    hasEOF = false
                                    break
                                }
                            }
                        }
                        if (hasEOF) {
                            add(EOF)
                        }
                    }
                }
            }
        }
        firstMap[symbol] = result
        result
    }

private fun followAll(
    symbol: NonTerminator, prods: ArrayList<Production>,
    start: NonTerminator,
    followAllMap: HashMap<NonTerminator, HashSet<Terminator>>,
    firstMap: HashMap<NonTerminator, HashSet<Terminator>>,
    // follow states map when last invoke this function
    tempFollowAllMap: HashMap<NonTerminator, HashSet<Terminator>> = hashMapOf(),
    // follow states size when last invoke this function
    // if size is same as last, states in
    lastVisitedFollowAllSize: HashMap<NonTerminator, Int> = hashMapOf(),
): HashSet<Terminator> =
    followAllMap[symbol] ?: run {
        val tempFollowSet = tempFollowAllMap[symbol]
        if (tempFollowSet != null && tempFollowSet.size == lastVisitedFollowAllSize[symbol]) {
            followAllMap[symbol] = tempFollowSet
            return@run tempFollowSet
        }
        lastVisitedFollowAllSize.putAll(tempFollowAllMap.mapValues { it.value.size })
        val result = HashSet<Terminator>().apply {
            if (symbol == start) {
                add(EOF)
            }
            val requiredFollowFromOtherNonTerminatorSet = HashSet<NonTerminator>()
            for (prod in prods) {
                for ((index, symbolInRight) in prod.right.withIndex()) {
                    if (symbolInRight != symbol) continue
                    if (index == prod.right.size - 1) {
                        if (prod.left != symbol) {
                            requiredFollowFromOtherNonTerminatorSet.add(prod.left)
                        }
                    } else if (prod[index + 1].isTerminator) {
                        val terminator = prod[index + 1] as Terminator
                        tempFollowAllMap.getOrPut(symbol) { hashSetOf() }.add(terminator)
                        add(terminator)
                    } else {
                        val terminators = first(prod[index + 1] as NonTerminator, prods, firstMap)
                        tempFollowAllMap.getOrPut(symbol) { hashSetOf() }.addAll(terminators)
                        addAll(terminators)
                    }
                }
            }
            for (otherNonTerminator in requiredFollowFromOtherNonTerminatorSet) {
                addAll(followAll(otherNonTerminator, prods, start, followAllMap, firstMap, tempFollowAllMap, lastVisitedFollowAllSize))
            }
        }
        followAllMap[symbol] = result
        result
    }

/**
 * 非终结符在特定产生式内的follow集
 */
private fun follow(
    index: Pair<Int, Int>, prods: ArrayList<Production>,
    start: NonTerminator,
    followMap: HashMap<Pair<Int, Int>, HashSet<Terminator>>,
    followAllMap: HashMap<NonTerminator, HashSet<Terminator>>,
    firstMap: HashMap<NonTerminator, HashSet<Terminator>>,
): HashSet<Terminator> =
    followMap[index] ?: run {
        val result = if (index.first == -1) {
            HashSet<Terminator>().apply { add(EOF) }
        } else if (prods[index.first].right.size > index.second + 1) {
            if (prods[index.first][index.second + 1].isTerminator) {//后跟的文法符号为终结符,直接加入
                HashSet<Terminator>().apply { add(prods[index.first][index.second + 1] as Terminator) }
            } else {
                first(
                    prods[index.first][index.second + 1] as NonTerminator,
                    prods, firstMap
                )
            }
        } else {
            followAll(prods[index.first].left, prods, start, followAllMap, firstMap)
        }
        followMap[index] = result
        result
    }

private fun epsilonAction(
    group: ProdGroup, prods: ArrayList<Production>,
    startProduction: Production, start: NonTerminator,
    followMap: HashMap<Pair<Int, Int>, HashSet<Terminator>>,
    followAllMap: HashMap<NonTerminator, HashSet<Terminator>>,
    firstMap: HashMap<NonTerminator, HashSet<Terminator>>,
): ProdGroup {
    val needVisit = ArrayDeque(group.prods)
    val visited = HashSet<Pair<Int, Int>>()
    val it = needVisit.iterator()
    while (it.hasNext()) {
        val nextProd = it.next()
        it.remove()
        if (visited.contains(nextProd)) continue
        visited.add(nextProd)
        val prodNow = if (nextProd.first != -1) prods[nextProd.first] else startProduction
        if (prodNow.right.isNotEmpty() && nextProd.second < prodNow.right.size) {
            if (!prodNow.right[nextProd.second].isTerminator) {//规约标记位于非终结符前面
                for ((pI, prod) in prods.withIndex()) {
                    if (prod.left == prodNow.right[nextProd.second]) {
                        val pairNext = Pair(pI, 0)
                        needVisit.add(pairNext)
                        group.reduceMap.getOrPut(pairNext) { HashSet() }
                            .addAll(follow(nextProd, prods, start, followMap, followAllMap, firstMap))
                    }
                }
            }
        }
    }
    return group
}

private fun symbolsCanMove(
    group: ProdGroup, prods: ArrayList<Production>, start: NonTerminator,
): HashSet<ISymbol> =
    HashSet<ISymbol>().apply {
        for (prodIndex in group.reduceMap.keys) {
            if (prodIndex.first == -1) {
                if (prodIndex.second == 0)
                    add(start)
                continue
            }
            if (prods[prodIndex.first].right.size > prodIndex.second) {//非规约项，规约标记不位于最后
                add(prods[prodIndex])
            }
        }
    }

private fun moveAction(
    group: ProdGroup, prods: ArrayList<Production>,
    start: NonTerminator, moveSymbol: ISymbol,
) = ProdGroup(
    HashMap<Pair<Int, Int>, HashSet<Terminator>>().apply {
        putAll(group.reduceMap.filterKeys {
            if (it.first == -1) {
                return@filterKeys moveSymbol == start && it.second == 0
            }
            val prod = prods[it.first]
            if (prod.right.size > it.second) {
                prod.right[it.second] == moveSymbol
            } else {
                false
            }
        }.mapKeys {
            Pair(it.key.first, it.key.second + 1)
        })
    },
    -1
)

private fun Production.format(pushedPos: Int): String {
    return "${left.name} --> ${
        right.withIndex()
            .joinToString { if (it.index == pushedPos) " [-] ${it.value.simpleString()}" else it.value.simpleString() }
    }"
}

@TestOnly
@Suppress("Unused")
private fun ProdGroup.format(prods: ArrayList<Production>): String {
    val sb = StringBuilder("[$state]: ")
    for (it in reduceMap) {
        sb.append("${if (it.key.first == -1) "S' --> S" else prods[it.key.first].format(it.key.second)}  ${it.value.map { it1 -> it1.simpleString() }}")
        sb.append("\n")
    }
    return sb.toString()
}

/**
 * Generate predict table
 *
 * @return Predict Table
 */
fun Grammar2.generatePredictTable(): PredictTable {
    val startProd = Production(EX_START_NON_TERMINATOR, ArrayList<ISymbol>().apply { add(start) })
    val exStartProd = Pair(-1, 0)
    val startReduceMap = HashMap<Pair<Int, Int>, HashSet<Terminator>>().apply {
        put(exStartProd, HashSet<Terminator>().apply { add(EOF) })
    }
    var state = 0
    val startGroup = ProdGroup(startReduceMap, state++)
    val prodGroups = HashSet<ProdGroup>().apply { add(startGroup) }
    val followAllMap = HashMap<NonTerminator, HashSet<Terminator>>().apply {
        put(EX_START_NON_TERMINATOR, HashSet<Terminator>().apply { add(EOF) })
    }
    val followMap = HashMap<Pair<Int, Int>, HashSet<Terminator>>()
    val firstMap = HashMap<NonTerminator, HashSet<Terminator>>()
    fun epsilonAction(group: ProdGroup) =
        epsilonAction(group, productions, startProd, start, followMap, followAllMap, firstMap)
    epsilonAction(startGroup)
    val needVisit = ArrayDeque<ProdGroup>().apply { add(startGroup) }
    val it = needVisit.iterator()
    val predictMap = HashMap<Int, HashMap<FAEdge<ISymbol>, Int>>()
    while (it.hasNext()) {
        val pGroupNow = it.next()
        it.remove()
        val moveSymbols = symbolsCanMove(pGroupNow, productions, start)
        for (symbol in moveSymbols) {
            val group = epsilonAction(moveAction(pGroupNow, productions, start, symbol))
            val nextState: Int
            val findGroup = prodGroups.find { it == group }
            if (findGroup != null) {
                nextState = findGroup.state
            } else {
                nextState = state++
                group.state = nextState
                prodGroups.add(group)
                needVisit.add(group)
            }
            if (symbol.isTerminator) {
                predictMap.getOrPut(pGroupNow.state) { HashMap() }[FAEdge(symbol, MOVE_IN)] = nextState
            } else {
                predictMap.getOrPut(pGroupNow.state) { HashMap() }[FAEdge(symbol, GOTO)] = nextState
            }
        }
        for ((pI, reduceSymbols) in pGroupNow.reduceMap) {
            if (pI.first == -1) {
                if (pI.second != 0) {
                    val nextStateMap = predictMap.getOrPut(pGroupNow.state) { HashMap() }
                    if (nextStateMap.any { it.key.value == EOF }) {
                        throw Exception("move in or reduce conflict!")
                    }
                    nextStateMap[FAEdge(EOF, REDUCE)] = -1
                }
                continue
            }
            if (productions[pI.first].right.size <= pI.second) {
                for (reduceSymbol in reduceSymbols) {
                    val nextStateMap = predictMap.getOrPut(pGroupNow.state) { HashMap() }
                    if (nextStateMap.any { it.key.value == reduceSymbol }) {
                        throw Exception("move in or reduce conflict in production ${productions[pI.first]}!")
                    }
                    nextStateMap[FAEdge(reduceSymbol, REDUCE)] = pI.first
                }
            }
        }
    }
    return PredictTable(productions, DFA.fromStateMap(0, HashSet(), predictMap))
}

/**
 * Generate sem predict table
 *
 * @return Semantic predict table
 */
fun SemGrammar.generateSemPredictTable(): SemPredictTable {
    val pt = generatePredictTable()
    return SemPredictTable(pt.productions, pt.dfa, semantics, globalDef)
}